from math import nan
import akshare as ak
import pandas as pd
from datetime import datetime
import time

from config.database import engine
from stockbase import *


def main():

    # 获取所有A股股票列表
    stock_list = fetch_all_a_stocks()
    if stock_list.empty:
        return

    # 保存股票基本信息
    stock_list.rename(columns={"code": "symbol"}, inplace=True)
    save_stock_basic(stock_list)

    # 获取股票基本信息
    stock_list = get_stock_basic()
    # stock_list = fetch_stock_not_synchr('monthly')
    # 限制获取的股票数量，测试时可以设置较小的值
    max_stocks = None  # 设为None获取全部股票
    symbols = stock_list['symbol'].head(
        max_stocks) if max_stocks else stock_list['symbol']  # type: ignore

    start_date = '20100101'
    for symbol in symbols:
        print(f"\n正在处理股票: {symbol} 的日线数据")

        # 获取日线数据
        period = 'daily'
        daily_data = fetch_daily_data(symbol, start_date, period)
        rename_data(daily_data)
        # amplitude振幅数据可能不对, 统一设置为0
        daily_data.loc[(daily_data['amplitude'] > 100000) | (
            daily_data['amplitude'] < -100000), 'amplitude'] = 0

        if not daily_data.empty:
            # 保存日线数据
            save_to_database(daily_data, 'stock_daily')

        # 避免请求过于频繁
        time.sleep(1)

        # 获取周线数据
        period = 'weekly'
        weekly_data = fetch_daily_data(symbol, start_date, period)
        rename_data(weekly_data)
        weekly_data.loc[(weekly_data['amplitude'] > 100000) | (
            weekly_data['amplitude'] < -100000), 'amplitude'] = 0
        if not weekly_data.empty:
            # 保存周线数据
            save_to_database(weekly_data, 'stock_weekly')

        # 避免请求过于频繁
        time.sleep(1)

        # 获取月线数据
        period = 'monthly'
        monthly_data = fetch_daily_data(symbol, start_date, period)
        rename_data(monthly_data)
        monthly_data.loc[(monthly_data['amplitude'] > 100000) | (
            monthly_data['amplitude'] < -100000), 'amplitude'] = 0
        if not monthly_data.empty:
            # 保存月线数据
            save_to_database(monthly_data, 'stock_monthly')

        # 避免请求过于频繁
        time.sleep(1)


if __name__ == "__main__":
    main()
